package com.atguigu.common.base.repository.impl;

import com.atguigu.common.base.entity.AutoIdBase;
import com.atguigu.common.base.repository.ParallelSaveRepository;
import org.springframework.data.jpa.repository.support.JpaEntityInformation;
import org.springframework.data.jpa.repository.support.SimpleJpaRepository;

import javax.persistence.EntityManager;
import javax.persistence.EntityManagerFactory;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.*;
import java.util.stream.Collectors;

public class ParallelSaveRepositoryImpl<T extends AutoIdBase, ID> extends SimpleJpaRepository<T, ID> implements ParallelSaveRepository<T, ID> {
    private final EntityManagerFactory entityManagerFactory;
    private final EntityManager entityManager;
    private final ThreadLocal<EntityManager> threadLocal = new ThreadLocal<>();


    public ParallelSaveRepositoryImpl(JpaEntityInformation<T, ?> entityInformation, EntityManager entityManager) {
        super(entityInformation, entityManager);
        this.entityManager = entityManager;
        this.entityManagerFactory = entityManager.getEntityManagerFactory();
    }

    public void parallelSave(List<T> inputs) {
        List<T> entities = inputs.stream().filter(AutoIdBase::isModified).collect(Collectors.toList());
        ExecutorService executorService = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors());

        List<EntitySplitWriter> writers = new ArrayList<>();
        int share = 500;
        for (int i = 0; i * share < entities.size(); i++) {
            List<T> entitySplit = entities.subList(i * share, Integer.min(i * share + share, entities.size()));
            writers.add(new EntitySplitWriter(entitySplit));
        }
        try {
            List<Future<Void>> futures = executorService.invokeAll(writers);

            for (Future<Void> future : futures) {
                future.get();
            }
        } catch (InterruptedException | ExecutionException e) {
            throw new RuntimeException(e);
        }
        executorService.shutdown();
    }

    @Override
    public <S extends T> S save(S entity) {
        if (entity.isModified()) {
            if (entity.isNew()) {
                entityManager.persist(entity);
                entity.setNew(false);
            } else {
                entityManager.merge(entity);
            }
            entity.setModified(false);
        }
        return entity;
    }

    public class EntitySplitWriter implements Callable<Void> {
        private final List<T> entitySplit;

        public EntitySplitWriter(List<T> entitySplit) {
            this.entitySplit = entitySplit;
        }

        @Override
        public Void call() {
            EntityManager entityManager = threadLocal.get();
            if (entityManager == null || !entityManager.isOpen()) {
                entityManager = entityManagerFactory.createEntityManager();
                threadLocal.set(entityManager);
            }
            try {
                entityManager.getTransaction().begin();
                for (T entity : entitySplit) {
                    if (entity.isNew()) {
                        entityManager.persist(entity);
                        entity.setNew(false);
                    } else {
                        entityManager.merge(entity);
                    }
                    entity.setModified(false);
                }
                entityManager.getTransaction().commit();
            } catch (Exception e) {
                entityManager.getTransaction().rollback();
                throw new RuntimeException(e);
            } finally {
                threadLocal.set(null);
                if (entityManager != null) {
                    entityManager.close();
                }
            }
            return null;
        }
    }

}
